/*
    Copyright 2006-2012 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Contains the dust model class, which holds a collection of scatterers.

#ifndef __dust_model__
#define __dust_model__

#include <vector> 
#include "boost/shared_ptr.hpp"
#include "mcrx-types.h"
#include "ray.h"
#include "random.h"
#include "constants.h"
#include "angular.h"
#include "chromatic_policy.h"
#include "boost/thread/mutex.hpp"

namespace mcrx {
  template <typename, typename> class generic_dust_model;
  template <typename, template<class> class, typename> class dust_model;
  template <typename> class intensity;
  template <typename, typename> class intensity_pointer;

  // typedef of commonly used polychromatic dust model type
  template<template<typename> class, typename> class scatterer;
  template<typename> class cumulative_sampling;
  typedef mcrx::dust_model<mcrx::scatterer<mcrx::polychromatic_scatterer_policy,
					   mcrx::mcrx_rng_policy>,
			   mcrx::cumulative_sampling,
			   mcrx::mcrx_rng_policy> T_polychromatic_dust_model;
}


/** A dust model class for the generic runs without scattering. It
    only defines the types for the xfer class. Note that this should
    not be used for array_1 shooting, because it doesn't size the
    arrays properly. For that, the full dust model with polychromatic
    policy should be used, even though it's more painful. */
template <typename content_type, typename rng_policy_type>
class mcrx::generic_dust_model {
public:
  typedef content_type T_lambda;
  typedef rng_policy_type T_rng_policy;
  typedef dummy_scatterer<T_lambda, T_rng_policy> T_scatterer;
  typedef generic_chromatic_policy<T_lambda> T_chromatic_policy;
  typedef typename T_chromatic_policy::T_biaser T_biaser;
  typedef ray<T_lambda> T_ray;

  T_densities zero_density() const {return T_densities();};
  T_lambda zero_lambda() const {return 0.;};

  // no-op implementations of all the dust_model members
  void set_wavelength(const T_lambda& lambda) {assert(0); };
  int n_scatterers() const {return scatterers_.size();};
  const T_scatterer get_scatterer (int i) const { 
    assert(0); return T_scatterer(); };
  T_scatterer get_scatterer (int i) {
    assert(0); return T_scatterer(); };
  
  array_1 reference_abs_lengths(const T_densities& rho, 
				const T_biaser& b) const {
    assert(0); return array_1(); };
  
  T_float total_reference_abs_length(const T_densities& rho, 
				     const T_biaser& b) const {
    assert(0); return 0; };
  
  T_lambda total_abs_length(const T_densities& rho) const {
    assert(0); return T_lambda(); };

  void 
  total_abs_length(const T_densities& rho, T_lambda& kappa_out) const {
    assert(0); };
  template <typename T> T_float
  total_abs_length(const blitz::ETBase<T>& rho) const {
    assert(0); return 0; };
  T_lambda total_abs_length_to_absorption(const T_densities& rho) const {
    assert(0); return 0; };
  T_lambda lambda () const {assert(0); return T_lambda(); };

  template <typename T_rng_policy>
  const angular_distribution<T_lambda> & scatter (const T_densities& rho,
						  T_biaser b,
						  T_ray& r, 
						  T_rng_policy& rng,
						  T_float& prob) const {
    assert(0); return isotropic_angular_distribution<T_lambda>(zero_lambda()); };
  template <typename T_rng_policy>
  int draw_scatterer (const T_densities& rho, T_biaser b,
		      T_ray& r, T_rng_policy& rng, T_float& prob) const {
    assert(0); return 0; };
  void add_absorbed(const T_lambda& i) const {
    assert(0); };
  const T_lambda& absorbed() const {assert(0); return T_lambda(); };
  void reset() {assert(0); };

};


/** The dust model class holds a collection of scatterers and can draw
    which scatterer will scatter a ray, depending on the local
    densities of the scattering species. */
template <typename scatterer_type,
	  template<class> class T_sampling_policy,
	  typename rng_policy_type = mcrx::global_random> 
class mcrx::dust_model {
public:
  typedef scatterer_type T_scatterer;
  typedef rng_policy_type T_rng_policy;
  typedef typename T_scatterer::T_chromatic_policy::T_dust_model_policy T_chromatic_policy;
  //typedef chromatic_policy<dust_model> T_chromatic_policy;
  typedef typename T_chromatic_policy::T_lambda T_lambda;
  //typedef typename T_chromatic_policy::T_scatterer T_scatterer;
  typedef typename T_chromatic_policy::T_scatterer_vector T_scatterer_vector;
  typedef typename T_chromatic_policy::T_opacity_array T_opacity_array;
  typedef typename T_chromatic_policy::T_biaser T_biaser;
  typedef ray<T_lambda> T_ray;

private:

  /// The wavelength vector used
  T_lambda lambda_;

  T_scatterer_vector scatterers_;
  /// The array of opacities. Indices are (scatterer, lambda).
  T_opacity_array opacities_;
  /** The opacity reference array. This array shares data with
      opacities_, but is not reference counted and will thus avoid
      mutex locks when copied. This is important for performance. */
  T_opacity_array opacities_ref_;

  mutable T_sampling_policy<T_rng_policy> distr_;

  mutable boost::mutex mutex_;

  /** Returns a copy of the opacities array which is not reference
      counted and bypasses mutex locking. This is safe only if you can
      guarantee that opacities_ is not deleted while this array exists
      (and that the array is not changed.) Should not be a problem
      inside the shoot loop, and it's necessary for performance. */
  const T_opacity_array& opacities_unsafe() const {
    return opacities_ref_;};

  /** Luminosity which has been absorbed by the dust grains (for
      debugging purposes). */
  mutable T_lambda absorbed_; 

public:
  dust_model() {};
  dust_model(const T_scatterer_vector& v) : scatterers_(v) {
    assert (v.size()> 0);};

  void set_wavelength(const T_lambda& lambda);
  int n_scatterers() const {return scatterers_.size();};
  const T_scatterer& get_scatterer (int i) const {
    assert(i>=0); assert(i<scatterers_.size()); return *scatterers_[i];};
  T_scatterer& get_scatterer (int i) {
    assert(i>=0); assert(i<scatterers_.size()); return *scatterers_[i];};

  /** Return an array of absorption lengths (which is dtau/dl) at the
      reference wavelength by multiplying the opacity array at the
      reference wavelength with a density array. */
  //typename T_chromatic_policy::RT_reference_abs_lengths
  array_1
  reference_abs_lengths(const T_densities& rho, 
			const T_biaser& b) const {
    /* BADNESS! The array returned by b.reference2 goes out of scope
       before the expression is evaluated! (Even thought no array
       should be constructed if optimization goes as planned. I think
       we need to return some expression thing that references the
       original array if we are to use ETs here. */
    array_1 ret(rho*b.reference2(opacities_unsafe()));
    threadLocal_warn (ret, true);
    return ret;
  };

  /** Return the total absorption length (which is dtau/dl) at the
      reference wavelength by summing the reference absorption length
      array over species. This could call reference_abs_lengths but
      that would involve making a temporary array and allocating
      memory so we stack it all into one expression. */
  T_float total_reference_abs_length(const T_densities& rho, 
				     const T_biaser& b) const {
    if(rho.size()==1)
      // special code path for one-species only bypasses sum
      return rho(0)*b.reference2_comp0(opacities_unsafe());
    else
      return sum(rho*b.reference2(opacities_unsafe()));
  };

  /** Return the total absorption length array (at all wavelengths) by
      summing the absorption length array over species. This just maps
      rho to dimension 2 and forwards to the chromatic_policy class
      for the actual calculation.  */
  T_lambda total_abs_length(const T_densities& rho) const {
    return T_chromatic_policy::total_abs_length
      (T_chromatic_policy::abs_length_array(rho(blitz::tensor::j), 
					    opacities_unsafe())); };

  /** Calculates the total absorption length array (at all
      wavelengths) by summing the absorption length array over
      species. It just forwards to the chromatic policy.  */
  inline void 
  total_abs_length(const T_densities& rho, T_lambda& kappa_out) const {
    return T_chromatic_policy::total_abs_length
      (rho, opacities_unsafe(), kappa_out); };

  /** Return the total absorption length array (at all wavelengths) by
      summing the absorption length array over species. This version
      takes a blitz expression as density vector. This just forwards
      to the chromatic_policy class for the actual calculation, so the
      same NOTE applies that rho has to be mapped to dimension 2, ie
      passed as rho(tensor::j).  */
  template <typename T>
  typename T_chromatic_policy::template RT_total_abs_length2<T>::T_result
  total_abs_length(const blitz::ETBase<T>& rho) const {
    return T_chromatic_policy::total_abs_length(rho, opacities_unsafe());};
    
  /** Return the total absorption length ("absorption opacity") to
      absorption only. This is used to calculate energy deposition
      rate from mean intensity. */
  T_lambda total_abs_length_to_absorption(const T_densities& rho) const;

  /** Returns a (thread-local) density vector of 0.  This is necessary
      because you must be able to initialize density vectors to
      appropriate size.  */
  T_densities zero_density () const {
    T_densities rho(scatterers_.size ());
    threadLocal_warn (rho, true);
    rho=0; return rho;};
  /** Returns a lambda vector of 0. Is guaranteed to not be shared, so
      can be copied by reference safely. */
  T_lambda zero_lambda () const {return get_scatterer(0).zero_lambda();};
  const T_lambda& lambda () const {return lambda_; };

  virtual const angular_distribution<T_lambda> & scatter (const T_densities& rho,
							  T_biaser b,
							  T_ray& r, 
							  T_rng_policy& rng,
							  T_float& prob) const;
  virtual int draw_scatterer (const T_densities& rho,
			      T_biaser b,
			      T_ray& r, 
			      T_rng_policy& rng, T_float& prob) const;

  /** Adds the specified luminosity to the absorbed luminosity
      accumulator. This requires a global mutex lock, but is only done
      for debugging. */
  void add_absorbed(const T_lambda& i) const {
    boost::mutex::scoped_lock lock(mutex_);absorbed_ += i;};
  const T_lambda& absorbed() const {return absorbed_;};

  /// Resets the absorbed luminosity accumulator.
  void reset() {absorbed_ = 0;};
};


// ***function definitions ***


/** Sets up the scatterers for the specified wavelength(s). The
    chromatic_policy classes supply the prefix/postfix functions
    which make sure things are set up correctly. */
template <typename scatterer_type,
	  template<class> class T_sampling_policy,
	  typename rng_policy_type>
void 
mcrx::dust_model<scatterer_type,
		 T_sampling_policy, 
		 rng_policy_type>::set_wavelength(const T_lambda& lambda)
{
  assign(lambda_, lambda);

  T_chromatic_policy::set_wavelength_prefix(opacities_, scatterers_, lambda);

  for (typename T_scatterer_vector::iterator i = scatterers_.begin();
       i != scatterers_.end(); ++i)
    (*i)->set_wavelength(lambda);

  T_chromatic_policy::set_wavelength_postfix(opacities_, scatterers_, lambda);

  // reset the ref array
  opacities_ref_.weakReference(opacities_);

  // zero the absorbed luminosity array, after making sure it's sized correctly
  resize_like (absorbed_, lambda);
  absorbed_=0;
}


/** Scatters a ray. The densities of the different species is used to
    draw which species will do the scattering. The probability of
    scattering off of different species changes with wavelength, so
    this introduces a possible bias factor. The scattering is then
    done by the scatterer for the species selected. */
template <typename scatterer_type,
	  template<class> class T_sampling_policy,
	  typename rng_policy_type>
const mcrx::angular_distribution<typename mcrx::dust_model<scatterer_type,
							   T_sampling_policy, 
							   rng_policy_type>::T_lambda>& 
mcrx::dust_model<scatterer_type,
		 T_sampling_policy, 
		 rng_policy_type>::scatter (const T_densities& rho,
					 T_biaser b,
					 T_ray& r, 
					 rng_policy_type& rng, T_float& prob) const
{
  T_float temp_prob;
  const int i = draw_scatterer(rho, b, r, rng, prob);
  const mcrx::angular_distribution<typename mcrx::dust_model<scatterer_type,
							     T_sampling_policy, 
							     T_rng_policy>::T_lambda>& 
    retval (scatterers_[i]->scatter(b, r, rng, temp_prob));
  prob *=temp_prob;
  return retval;
}

/** Draws the scatterer that does the interaction based on their
    opacities. Applies bias factor from opacities as a function of
    wavelength. */
template <typename scatterer_type,
	  template<class> class T_sampling_policy,
	  typename rng_policy_type>
int mcrx::dust_model<scatterer_type, //chromatic_policy, 
		     T_sampling_policy, 
		     rng_policy_type>::draw_scatterer (const T_densities& rho,
						    T_biaser b,
						    T_ray& r, 
						    T_rng_policy& rng,
						    T_float& prob) const
{
  // If there is only one scatterer, this is a noop that just returns
  // 0. There is no biasing in that case.
  if(rho.size()==1) {
    prob=1;
    return 0;
  }

  // calculate mean free paths for all scatterers/wavelengths
  const T_opacity_array mfps
    (T_chromatic_policy::abs_length_array(rho(blitz::tensor::i), opacities_unsafe()));
  const T_densities ref_mfps = b.reference2(mfps);
  // CRITICAL SECTION!!!
  int i;
  {
    boost::mutex::scoped_lock lock(mutex_);
    distr_.set_weights(ref_mfps);
    i = distr_.draw_entry(rng, prob);
  }

  // bias factor is bias_factor(mfp_i)/bias_factor(total absorption length)
  DEBUG(1,T_lambda bias(b.bias_factor(T_chromatic_policy::extract_scatterer_component(mfps, i))/ \
			b.bias_factor(T_chromatic_policy::total_abs_length(mfps))); \
	if(!condition_all(bias<1e3)) \
	  std::cout << "WARNING: Large bias factor  " << max(bias)<< " at lambda " << maxIndex(bias) << " reflambda " << b.reference_index() << " at " << __FILE__ <<':' <<__LINE__ << std::endl;);

  r.bias
    (b.bias_factor(T_chromatic_policy::extract_scatterer_component(mfps, i))/
     b.bias_factor(T_chromatic_policy::total_abs_length(mfps)));

  return i;
}

template <typename scatterer_type,
	  template<class> class T_sampling_policy,
	  typename rng_policy_type>
typename mcrx::dust_model<scatterer_type,
			  T_sampling_policy, 
			  rng_policy_type>::T_lambda
mcrx::dust_model<scatterer_type,
		 T_sampling_policy, 
		 rng_policy_type>::
total_abs_length_to_absorption(const T_densities& rho) const
{
  // If we could return this as an expression it would be much faster.
  T_lambda ao(zero_lambda());
  int ii=0;
  for(typename T_scatterer_vector::const_iterator i=scatterers_.begin();
      i!=scatterers_.end(); ++i, ++ii) {
    assert(all(opacities_unsafe()(ii, blitz::Range::all())==(*i)->opacity()));
    ao += rho(ii) * opacities_unsafe()(ii, blitz::Range::all()) * 
      (1 - (*i)->albedo());
  }
  return ao;
}

#endif


